from typing import List, Iterable

import cv2
import numpy as np
import torch
import nibabel

from config import image_root, label_seg_root, label_cls_path, batch_size
from .interface import Dataset as DatasetInterface


class Dataset(DatasetInterface):
    def __init__(self, names: List[str]):
        with open(label_cls_path, 'r') as f:
            f.readline()
            label_cls = f.readlines()
            label_cls = [line.split(',') for line in label_cls]
            label_cls = {name: [liver, spleen, left_kidney, right_kidney] for name, liver, spleen, left_kidney, right_kidney in label_cls}
        self.names = names
        self.source = {name: {
            'image': nibabel.load(image_root / f'{name}.nii.gz').get_fdata().astype(np.int16),
            'label_seg': nibabel.load(label_seg_root / f'{name}_mask.nii.gz').get_fdata().astype(np.uint8),
            'label_cls': np.array(label_cls[name]).astype(np.bool),
        } for name in names}
        # for name, item in self.source.items():
        #     item['image'] = cv2.resize(item['image'], dsize=(128, 128), interpolation=cv2.INTER_LINEAR)[:, :, :20]
        #     item['label_seg'] = cv2.resize(item['label_seg'], dsize=(128, 128), interpolation=cv2.INTER_NEAREST)[:, :, :20]
        self._mode = 'train_seg'
        self._sample = [(name, channel) for name in self.source.keys() for channel in range(self.source[name]['image'].shape[2])]

    def mode(self, mode: str = None) -> str:
        if mode and mode == 'train_seg':
            # 训练分割的时候，解锁分割网络梯度
            self._mode = mode
        elif mode and mode == 'train_cls':
            # 训练分类的时候，锁定分割网络梯度
            self._mode = mode
        elif mode and mode == 'evaluate':
            self._mode = mode
        return self._mode

    def sample(self) -> None:
        if self._mode == 'train_seg':
            self._sample = [(name, channel) for name in self.names for channel in range(self.source[name]['image'].shape[2])]
            self._sample += [self._sample[-1]] * (-len(self._sample) % batch_size)
            np.random.shuffle(self._sample)
        elif self._mode == 'train_cls':
            self._sample = [name for name in self.names]
            np.random.shuffle(self._sample)
        elif self._mode == 'evaluate':
            self._sample = [name for name in self.names]

    def __len__(self) -> int:
        return len(self._sample)

    def __getitem__(self, index) -> Iterable[np.ndarray]:
        if self._mode == 'train_seg':
            name, channel = self._sample[index]
            image = self.source[name]['image'][:, :, channel].astype(np.float32) / 4096
            label = self.source[name]['label_seg'][:, :, channel]
            label = np.eye(5, dtype=np.bool)[label].astype(np.float32)
            return image, label
        if self._mode == 'train_cls':
            name = self._sample[index]
            image = self.source[name]['image'].astype(np.float32) / 4096
            label = self.source[name]['label_cls'].astype(np.float32)
            # !!!!!!!!!!!!!!!!!!!!!!!!
            image = cv2.resize(image, (128, 128))
            return image, label
        if self._mode == 'evaluate':
            name = self._sample[index]
            image = self.source[name]['image'].astype(np.float32) / 4096
            label_seg = self.source[name]['label_seg']
            label_seg = np.eye(5, dtype=np.bool)[label_seg].astype(np.float32)
            label_cls = self.source[name]['label_cls'].astype(np.float32)
            # !!!!!!!!!!!!!!!!!!!!!!!!
            image = cv2.resize(image, (128, 128))
            label_seg = np.zeros(shape=(*image.shape, 5), dtype=np.float32)
            return image, label_seg, label_cls
        raise Exception(f'mode not allowed {self._mode}')

    def __iter__(self) -> Iterable[np.ndarray]:
        if self._mode == 'train_seg':
            for x in range(0, len(self), batch_size):
                images, labels = zip(*[self[x + i] for i in range(batch_size)])
                images = torch.tensor(images, dtype=torch.float32).permute(0, 2, 1)[:, None, :, :]
                labels = torch.tensor(labels, dtype=torch.float32).permute(0, 3, 2, 1)
                yield images, labels
        if self._mode == 'train_cls':
            for x in range(len(self)):
                image, label = self[x]
                image = torch.tensor(image, dtype=torch.float32).permute(2, 1, 0)[:, None, :, :]
                label = torch.tensor(label, dtype=torch.float32)[None, :]
                yield image, label
        if self._mode == 'evaluate':
            for x in range(len(self)):
                image, label_seg, label_cls = self[x]
                image = torch.tensor(image, dtype=torch.float32).permute(2, 1, 0)[:, None, :, :]
                label_seg = torch.tensor(label_seg, dtype=torch.float32).permute(2, 3, 1, 0)
                label_cls = torch.tensor(label_cls, dtype=torch.float32)[None, :]
                yield image, label_seg, label_cls
